{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir(os.path.split(os.getcwd())[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import gym\n",
    "from agent import *\n",
    "from optionpricing import *\n",
    "import yaml\n",
    "import torch\n",
    "from collections import defaultdict\n",
    "import matplotlib.style as style\n",
    "import ipywidgets as widgets\n",
    "from ipywidgets import interact, interact_manual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "style.use('seaborn-poster')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_folder = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join('experiments', experiment_folder, 'config.yaml'), 'r') as f:\n",
    "    args_dict = yaml.load(f, Loader = yaml.SafeLoader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    def __init__(self, **kwargs):\n",
    "        self.__dict__.update(kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = Args(**args_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "        'S': 100,\n",
    "        'T': 10, # 10 days\n",
    "        'L': 1,\n",
    "        'm': 100, # L options for m stocks\n",
    "        'n': 0,\n",
    "        'K': [95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105],\n",
    "        'D': 5,\n",
    "        'mu': 0,\n",
    "        'sigma': 0.01,\n",
    "        'r': 0,\n",
    "        'ss': 5,\n",
    "        'kappa': 0.1,\n",
    "        'multiplier': args.trc_multiplier,\n",
    "        'ticksize': args.trc_ticksize,\n",
    "        }\n",
    "\n",
    "env = OptionPricingEnv(config)\n",
    "env.configure()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "ngpu = 1 if torch.cuda.is_available() else 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_estimator(env, device, nhidden, nunits, experiment_folder, kind = 'best'):\n",
    "    state_shape = env.observation_space.shape\n",
    "    state_space_dim = state_shape[0] if len(state_shape) == 1 else state_shape\n",
    "    \n",
    "    estimator = Estimator(nhidden, nunits, state_space_dim, env.action_space.n)\n",
    "    if kind == 'best':\n",
    "        checkpoint = torch.load(os.path.join('experiments', experiment_folder, 'best.pt'), map_location = torch.device('cpu'))\n",
    "    elif kind == 'checkpoint':\n",
    "        checkpoint = torch.load(os.path.join('experiments', experiment_folder, 'checkpoint.pt'), map_location = torch.device('cpu'))\n",
    "    else:\n",
    "        raise ValueError('Invalid choice for kind')\n",
    "        \n",
    "    estimator.load_state_dict(checkpoint['estimator'])\n",
    "    estimator.eval()\n",
    "    \n",
    "    return estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_episode(env, device, estimator, policy = 'agent', seed = 1):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    state = torch.from_numpy(env.reset()).to(device)\n",
    "    history = defaultdict(list)\n",
    "    done = False\n",
    "    \n",
    "    while not done:\n",
    "        history['delta'].append(env.delta)\n",
    "        if policy == 'agent':\n",
    "            with torch.no_grad():\n",
    "                action = np.argmax(estimator(state).numpy())\n",
    "        elif policy == 'delta':\n",
    "            action = delta_neutral_policy(env)\n",
    "            \n",
    "        state, reward, done, info = env.step(action)\n",
    "        \n",
    "        history['reward'].append(reward)\n",
    "        history['n'].append(env.n)\n",
    "        history['stock_value'].append(env.stock_value)\n",
    "        history['option_value'].append(env.option_value)\n",
    "        history['cash'].append(env.cash)\n",
    "        \n",
    "        state = torch.from_numpy(state).to(device)\n",
    "        \n",
    "    return history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delta_neutral_policy(env):\n",
    "    return env.inv_action_map[-1 * int(env.delta * (env.L * env.m)) - env.n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = load_estimator(env, device, args.nhidden, args.nunits, experiment_folder, 'best')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = simulate_episode(env, device, estimator, 'agent', seed = 7)\n",
    "\n",
    "stock_pnl = np.array(history['stock_value'][1:]) - np.array(history['stock_value'][:-1])\n",
    "stock_pnl = np.insert(stock_pnl, 0, 0)\n",
    "\n",
    "option_pnl = np.array(history['option_value'][1:]) + np.array(history['cash'][1:]) - np.array(history['option_value'][:-1]) - np.array(history['cash'][:-1])\n",
    "option_pnl = np.insert(option_pnl, 0, 0)\n",
    "\n",
    "total_pnl = stock_pnl + option_pnl\n",
    "\n",
    "steps = np.arange(1, len(history['delta']) + 1, 1)\n",
    "\n",
    "fig, ax = plt.subplots(figsize = (12, 8), nrows = 2, ncols = 1, sharex = True)\n",
    "ax[0].plot(steps, np.array(history['delta']) * -env.L * env.m, color = 'green', label = 'delta', lw = 1.5)\n",
    "ax[0].plot(steps, history['n'], color = 'blue', label = 'n', lw = 1.5)\n",
    "ax[0].legend()\n",
    "#ax[1].plot(steps, history['cash'], color = 'red', label = 'cash', lw = 1.5)\n",
    "ax[1].plot(steps, history['reward'], color = 'blue', label = 'reward', lw = 1.5)\n",
    "ax[1].plot(steps, np.clip(history['reward'], -args.clip, args.clip), color = 'red', label = 'clipped', lw = 1.5)\n",
    "#ax.plot(steps, total_pnl, lw = 1.5, label = 'total pnl', color = 'red')\n",
    "ax[0].set_title(f'K: {env.K}')\n",
    "ax[1].set_xlabel('step')\n",
    "ax[0].set_ylabel('n')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = simulate_episode(env, device, estimator, 'delta', seed = 7)\n",
    "\n",
    "stock_pnl = np.array(history['stock_value'][1:]) - np.array(history['stock_value'][:-1])\n",
    "stock_pnl = np.insert(stock_pnl, 0, 0)\n",
    "\n",
    "option_pnl = np.array(history['option_value'][1:]) + np.array(history['cash'][1:]) - np.array(history['option_value'][:-1]) - np.array(history['cash'][:-1])\n",
    "option_pnl = np.insert(option_pnl, 0, 0)\n",
    "\n",
    "total_pnl = stock_pnl + option_pnl\n",
    "\n",
    "steps = np.arange(1, len(history['delta']) + 1, 1)\n",
    "\n",
    "fig, ax = plt.subplots(figsize = (12, 8), nrows = 2, ncols = 1, sharex = True)\n",
    "ax[0].plot(steps, np.array(history['delta']) * -env.L * env.m, color = 'green', label = 'delta', lw = 1.5)\n",
    "ax[0].plot(steps, history['n'], color = 'blue', label = 'n', lw = 1.5)\n",
    "ax[0].legend()\n",
    "#ax[1].plot(steps, history['cash'], color = 'red', label = 'cash', lw = 1.5)\n",
    "ax[1].plot(steps, history['reward'], color = 'blue', label = 'reward', lw = 1.5)\n",
    "ax[1].plot(steps, np.clip(history['reward'], -args.clip, args.clip), color = 'red', label = 'clipped', lw = 1.5)\n",
    "#ax.plot(steps, total_pnl, lw = 1.5, label = 'total pnl', color = 'red')\n",
    "ax[0].set_title(f'K: {env.K}')\n",
    "ax[1].set_xlabel('step')\n",
    "ax[0].set_ylabel('n')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate State Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_state_space_data(env, device, estimator, data, variable = 'S'):\n",
    "    env.reset() # Only need S0 even though strike might be different\n",
    "    S0 = env.S0\n",
    "    \n",
    "    # Create input states\n",
    "    if variable == 'S':\n",
    "        S = data['S'].reshape(-1, 1) / S0\n",
    "        t = np.ones(S.shape) * data['t']\n",
    "        n = np.ones(S.shape) * data['n'] / env.high\n",
    "        K = np.ones(S.shape) * data['K'] / S0\n",
    "    \n",
    "    elif variable == 'n':\n",
    "        n = data['n'].reshape(-1, 1) / env.high\n",
    "        t = np.ones(n.shape) * data['t']\n",
    "        S = np.ones(n.shape) * data['S'] / S0\n",
    "        K = np.ones(n.shape) * data['K'] / S0\n",
    "    \n",
    "    elif variable == 't':\n",
    "        t = data['t'].reshape(-1, 1)\n",
    "        S = np.ones(t.shape) * data['S'] / S0\n",
    "        n = np.ones(t.shape) * data['n'] / env.high\n",
    "        K = np.ones(t.shape) * data['K'] / S0\n",
    "    \n",
    "    else:\n",
    "        raise ValueError\n",
    "        \n",
    "    state = np.hstack((S, t, n, K)).astype(np.float32)\n",
    "    state = torch.from_numpy(state).to(device)\n",
    "    with torch.no_grad():\n",
    "        action = np.argmax(estimator(state).numpy(), axis = 1)\n",
    "        \n",
    "    action = [env.action_map[a] for a in action]\n",
    "    \n",
    "    return action"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Action vs State Variable"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed S"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact\n",
    "def plot_S(n = (-100, 100, 1), t = (0, 1, 0.1), K = (90, 110, 1)):\n",
    "    S = np.linspace(85, 115, 10000)\n",
    "    fig, ax = plt.subplots(figsize = (12, 8))\n",
    "    actions = generate_state_space_data(env, device, estimator, {'S': S, 't': t, 'n': n, 'K': K}, variable = 'S')\n",
    "    ax.plot(S, actions, lw = 1.5, color = 'blue')\n",
    "    ax.set_xlabel('S')\n",
    "    ax.set_ylabel('action')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact\n",
    "def plot_n(S = (85, 115, 1), t = (0, 1, 0.1), K = (90, 110, 1)):\n",
    "    n = np.linspace(-100, 100, 10000)\n",
    "    fig, ax = plt.subplots(figsize = (12, 8))\n",
    "    actions = generate_state_space_data(env, device, estimator, {'S': S, 't': t, 'n': n, 'K': K}, variable = 'n')\n",
    "    ax.plot(n, actions, lw = 1.5, color = 'blue')\n",
    "    ax.set_xlabel('n')\n",
    "    ax.set_ylabel('action')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact\n",
    "def plot_t(S = (85, 115, 1), n = (-100, 100, 1), K = (90, 110, 1)):\n",
    "    t = np.linspace(0, 1, 10000)\n",
    "    fig, ax = plt.subplots(figsize = (12, 8))\n",
    "    actions = generate_state_space_data(env, device, estimator, {'S': S, 't': t, 'n': n, 'K': K}, variable = 't')\n",
    "    ax.plot(t, actions, lw = 1.5, color = 'blue')\n",
    "    ax.set_xlabel('t')\n",
    "    ax.set_ylabel('action')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
